import json
import time
from datetime import datetime
from decimal import Decimal

import ccxt
from apscheduler.executors.asyncio import AsyncIOExecutor
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from ccxt import Exchange
from jqdatasdk import *

from loguru import logger
from sqlalchemy import desc
from sqlalchemy.orm import Session

from app.api import deps
from app.common.dateutils import DateUtils
from app.core.config import settings
from app.core.constant_code import ConstantItem, ConstantJQDataFrequency, ExchangeType, MarketType
from app.core.dingtalk import RobotFactory
from app.core.exception_code import ExceptionCode, BasicExceptionCode
from app.core.exceptions import BizException
from app.models.bar import Bar
from app.models.instrument import Instrument
from app.spider.spider import Spider
import pandas as pd


class JqdataSpier(Spider):
    def __init__(self):
        self.client = None

    def getClient(self):
        auth(settings.JQDATA_KEY, settings.JQDATA_SECRET)


class JqdataStockInstrumentSpider(JqdataSpier):
    def __init__(self, db):
        super(JqdataStockInstrumentSpider, self).__init__()
        self.db = db
        self.client = self.getClient()

    def crawl(self):
        current = None
        try:
            robot = RobotFactory.getAvailable()
            if robot:
                robot.sendText(f"[{datetime.now()}] 商品采集(JqdataStockInstrumentSpider),任务开始......")
            else:
                logger.info('无法向钉钉发送消息......')

            stocks = get_all_securities(types=['stock'])
            stocks['symbol'] = stocks.index
            totalTimeStart = time.time()
            for index, item in stocks.iterrows():
                current = item
                arr = item['symbol'].split(".")
                instrument = Instrument()
                instrument.base = arr[0]
                instrument.quote = arr[1]
                instrument.symbol = item['symbol']
                instrument.name = item['display_name']
                instrument.precision = json.dumps({"amount": 2, "price": 2})
                instrument.limits = json.dumps({"amount": {"min": 0, "max": 0}, "price": {"min": 0, "max": 0}})
                # instrument.precision = json.dumps(item['precision'])
                # instrument.limits = json.dumps(item['limits'])
                instrument.nameLetter = item['name']
                instrument.areaId = 2
                instrument.exchangeId = ExchangeType.CHINA_STOCK
                # instrument.area = row['area']
                # instrument.industry = row['industry']
                # instrument.symbolPartition = row['market']
                # instrument.pricePrecision = symbol['baseAssetPrecision']
                # instrument.amountPrecision = symbol['baseAssetPrecision']
                # instrument.marketType = ConstantItem.MARKET_STOCKS
                # instrument.nameLetter = symbol['symbol']
                # instrument.exchange_id = self.exchange
                # instrument.countryCode = self.countryCode

                dbInstrument = self.db.query(Instrument).filter_by(areaId=instrument.areaId,
                                                                   exchangeId=instrument.exchangeId,
                                                                   symbol=instrument.symbol).one_or_none()
                if dbInstrument is not None:
                    dbInstrument.name = instrument.name
                    dbInstrument.base = instrument.base
                    dbInstrument.quote = instrument.quote
                    dbInstrument.nameLetter = instrument.nameLetter
                    self.db.commit()
                else:
                    self.db.add(instrument)
                    self.db.commit()
            self.db.close()
            msg = f"[{datetime.now()}] 商品采集(JqdataStockInstrumentSpider)=> 总共处理{len(stocks)}个商品，耗时{time.time() - totalTimeStart}，成功{len(stocks)}个，失败0个，失败商品列表：[]"
            logger.info(msg)
            robot = RobotFactory.getAvailable()
            if robot:
                robot.sendText(msg)
        except Exception as ex:
            logger.error("商品采集(JqdataStockInstrumentSpider)任务异常：" + json.dumps(current))
            logger.exception(ex)
        finally:
            robot = RobotFactory.getAvailable()
            if robot:
                robot.sendText(f"[{datetime.now()}] 商品采集(JqdataStockInstrumentSpider),任务结束......")


class JqdataStockBarSpider(JqdataSpier):
    def __init__(self, db, frequency, startDate, instruments=[]):
        super().__init__()
        self.db = db
        self.exchangeId = ExchangeType.CHINA_STOCK
        self.limit = 1500
        self.frequency = frequency
        self.startDate = startDate
        self.instruments = instruments if len(instruments) > 0 else self.getInstruments()
        self.client: ccxt.binance = self.getClient()

    def getInstruments(self):
        return self.db.query(Instrument).filter_by(exchangeId=self.exchangeId).all()

    def formatField(self, origin, bar):
        # result = get_price(instrument.symbol, start_date=DateUtils.format(tmpDate, DateUtils.DATE_TIME),
        #                    end_date=DateUtils.format(tmpEndDate, DateUtils.DATE_TIME), frequency='1d',
        #                    fields=['open', 'close', 'low', 'high', 'volume', 'money', 'factor',
        #                            'high_limit', 'low_limit', 'avg', 'pre_close', 'paused'], skip_paused=False, fq='pre', count=None)

        bar.barDate = origin['bar_date']
        # fetch_ohlcv
        bar.open = Decimal(str(origin['open']))
        bar.high = Decimal(str(origin['high']))
        bar.low = Decimal(str(origin['low']))
        bar.close = Decimal(str(origin['close']))
        bar.volume = Decimal(str(origin['volume']))
        bar.amount = Decimal(str(origin['money']))
        bar.preClose = Decimal(str(origin['pre_close']))
        bar.tradeStatus = 1 if origin['paused'] == 0 else 0
        bar.factor = Decimal(str(origin['factor']))
        return bar


class DefaultJqdataStockBarSpider(JqdataStockBarSpider):
    def __init__(self, db, frequency, startDate, instruments=[]):
        super(DefaultJqdataStockBarSpider, self).__init__(db, frequency, startDate, instruments)

    def crawl(self):
        handleCount = 0
        total = len(self.instruments)
        totalTimeStart = time.time()
        limit = 1500
        errors = []
        try:
            robot = RobotFactory.getAvailable()
            if robot:
                robot.sendText(f"[{datetime.now()}] K线（{self.frequency}）采集(DefaultJqdataStockBarSpider),任务开始运行......")
            logger.info(f"[{datetime.now()}] K线（{self.frequency}）采集(DefaultJqdataStockBarSpider),任务开始运行......")
            for instrument in self.instruments:
                try:
                    handleCount += 1
                    logger.info('{}-{}，当前进度{}%,总{},已采集{}', instrument.name, self.frequency, round(handleCount * 100 / total, 2), total, handleCount)
                    # 先查出库中最后一条数据
                    dbBar = self.db.query(Bar).filter(Bar.instrumentId == instrument.id,
                                                      Bar.frequency == self.frequency,
                                                      Bar.adjustFlag == ConstantItem.ADJUST_NO).order_by(desc(Bar.barDate)).first()
                    startDate = dbBar.barDate if dbBar else self.startDate
                    endDate = DateUtils.stock_daily_end()
                    tmpDate = startDate
                    rs = []
                    while tmpDate < endDate:
                        tmpEndDate = DateUtils.addDays(tmpDate, limit)
                        if tmpEndDate > endDate: tmpEndDate = endDate
                            
                        # 不复权数据 factor为空，需要额外处理因子 'factor'
                        result = None
                        result_factor = None
                        if self.frequency == '1d':
                            start_dt = DateUtils.format(tmpDate, DateUtils.DATE_TIME)
                            end_dt = DateUtils.format(tmpEndDate, DateUtils.DATE_TIME)
                            # 取日线数据, 不复权不提供factor
                            result = get_price(instrument.symbol,
                                               start_date=start_dt,
                                               end_date=end_dt,
                                               frequency=self.frequency,
                                               fields=['open', 'close', 'low', 'high', 'volume', 'money', 'high_limit',
                                                       'low_limit', 'avg', 'pre_close', 'paused'],
                                               skip_paused=True,
                                               fq='none')
                            # 取复权因子
                            result_factor = get_price(instrument.symbol,
                                               start_date=start_dt,
                                               end_date=end_dt,
                                               frequency=self.frequency,
                                               fields=['factor'],
                                               skip_paused=True,
                                               fq='post')
                        else:
                            # 取非日线数据
                            result = get_price(instrument.symbol,
                                               start_date=DateUtils.format(tmpDate, DateUtils.DATE_TIME),
                                               end_date=DateUtils.format(tmpEndDate, DateUtils.DATE_TIME),
                                               frequency=self.frequency,
                                               fields=['open', 'close', 'low', 'high', 'volume', 'money'],
                                               skip_paused=True,
                                               fq='none')
                            
                            # 从数据库日线中取复权因子/涨停价/跌停价 factor, high_limit, low_limit
                            result_factor = self.db.query(Bar).filter(Bar.instrumentId == instrument.id,
                                                      Bar.frequency == '1d',
                                                      Bar.adjustFlag == ConstantItem.ADJUST_NO,
                                                      Bar.barDate >= tmpDate,
                                                      Bar.barDate <= tmpEndDate).order_by(desc(Bar.barDate)).all()
                            print("从数据库日线中取前复权因子", "-" * 20)
                            print("result_factor", result_factor)
                            pass
                            
                        result['factor'] = result_factor['factor'].values
                        result['bar_date'] = result.index
                        if len(result) > 0:
                            rs = pd.concat([rs, result], ignore_index=False) if len(rs) > 0 else result
                            rs = rs[~rs.index.duplicated(keep='first')]
                        logger.info('{}-{}已采集{}条', instrument.name, self.frequency, len(rs))
                        # 获取最后一条数据的时间
                        # if len(result) == 0:
                        #     break
                        tmpDate = DateUtils.addDays(tmpDate, limit)
                        
                    rs = rs.reset_index(drop=True)
                    start = time.time()
                    if len(rs) == 0:
                        logger.info('{}-{}没有数据', instrument.name, self.frequency)
                    logger.info('获取{}条{}-{}数据,耗时{}秒', len(rs), instrument.symbol, self.frequency, time.time() - start)
                    bars = []
                    # 处理数据耗时
                    start = time.time()
                    count = 0
                    preRow = dbBar
                    if len(rs) > 0:
                        # 处理第一条数据，第一条有数据，做更新，无数据，做插入
                        first = rs.iloc[0]
                        if preRow:
                            # 做更新
                            if first['bar_date'] == preRow.barDate:
                                preRow = self.formatField(first, preRow)
                                preRow.change = Decimal(str(first['close'])) - preRow.preClose
                                preRow.pctChange = round(preRow.change * 100 / preRow.preClose, 4)
                                # preRow.amount = 0
                            else:
                                preRow = self.formatField(first, Bar())
                                preRow.adjustFlag = ConstantItem.ADJUST_NO
                                preRow.frequency = self.frequency
                                preRow.instrumentId = instrument.id
                                preRow.tradeStatus = 1
                                # preRow.preClose = dbBar.close
                                preRow.change = preRow.close - preRow.preClose
                                preRow.pctChange = round(preRow.change * 100 / preRow.preClose, 4)
                                self.db.add(preRow)
                        else:
                            # 插入
                            preRow = self.formatField(first, Bar())
                            preRow.adjustFlag = ConstantItem.ADJUST_NO
                            preRow.frequency = self.frequency
                            preRow.instrumentId = instrument.id
                            # preRow.tradeStatus = 1
                            preRow.change = preRow.close - preRow.preClose
                            # preRow.preClose = first[4]
                            preRow.pctChange = round(preRow.change * 100 / preRow.preClose, 4)
                            # preRow.amount = 0

                            self.db.add(preRow)
                    self.db.commit()

                    for i, row in rs.iterrows():
                        if i == 0:
                            preRow = row
                            continue
                        bar = self.formatField(row, Bar())
                        bar.adjustFlag = ConstantItem.ADJUST_NO
                        bar.frequency = self.frequency
                        bar.instrumentId = instrument.id
                        bar.tradeStatus = 1
                        bar.change = bar.close - bar.preClose
                        # bar.preClose = Decimal(preRow[4])
                        # bar.amount = 0
                        bar.pctChange = round(bar.change * 100 / bar.preClose, 4)
                        bars.append(bar)
                        # 分段插入
                        if len(bars) == 200000 or i == len(rs) - 1:
                            count = count + len(bars)
                            logger.info('已处理{}条{}-{}数据耗时{}秒', count, instrument.symbol, self.frequency, time.time() - start)
                            self.db.bulk_save_objects(bars)
                            # self.db.flush()
                            self.db.commit()
                            bars.clear()
                        preRow = row
                    # self.db.flush()
                    self.db.commit()
                    logger.info('总共{}条({},{})-{}数据,耗时{}秒', len(rs), instrument.symbol, instrument.name, self.frequency,
                                time.time() - start)
                except Exception as ex:
                    msg = f'{instrument.symbol}-{self.frequency}'
                    logger.error(f"{msg}处理失败")
                    logger.exception(ex)
                    errors.append(msg)
            msg = f'[{datetime.now()}] K线（{self.frequency}）爬取(JqdataStockInstrumentSpider) => 总共处理{len(self.instruments)}个商品，耗时{time.time() - totalTimeStart}，成功{len(self.instruments) - len(errors)}个，失败{len(errors)}个，失败商品列表：[{errors}]'
            logger.info(msg)
            robot = RobotFactory.getAvailable()
            if robot:
                robot.sendText(msg)
        except Exception as ex:
            logger.info(f"[{datetime.now()}] K线采集(JqdataStockInstrumentSpider),采集失败......")
            logger.exception(ex)
        finally:
            logger.info(f"[{datetime.now()}] K线采集(JqdataStockInstrumentSpider),任务结束......")
            robot = RobotFactory.getAvailable()
            if robot:
                robot.sendText(f"[{datetime.now()}] K线采集(JqdataStockInstrumentSpider),任务结束......")


def getApiFrequency(frequency):
    # if frequency == ConstantItem.FREQUENCY_1S:
    #     return '1s'
    if frequency == ConstantItem.FREQUENCY_1MIN:
        return '1m'
    elif frequency == ConstantItem.FREQUENCY_3MIN:
        return '3m'
    elif frequency == ConstantItem.FREQUENCY_5MIN:
        return '5m'
    elif frequency == ConstantItem.FREQUENCY_15MIN:
        return '15m'
    elif frequency == ConstantItem.FREQUENCY_30MIN:
        return '30m'
    elif frequency == ConstantItem.FREQUENCY_1HOUR:
        return '1h'
    elif frequency == ConstantItem.FREQUENCY_2HOUR:
        return '2h'
    elif frequency == ConstantItem.FREQUENCY_4HOUR:
        return '4h'
    elif frequency == ConstantItem.FREQUENCY_6HOUR:
        return '6h'
    elif frequency == ConstantItem.FREQUENCY_8HOUR:
        return '8h'
    elif frequency == ConstantItem.FREQUENCY_12HOUR:
        return '12h'
    elif frequency == ConstantItem.FREQUENCY_1D:
        return '1d'
    elif frequency == ConstantItem.FREQUENCY_1W:
        return '1w'
    else:
        raise BizException(BasicExceptionCode.ILLEGAL_PARAMETER)


def frequencyUnit(frequency):
    if frequency == ConstantItem.FREQUENCY_1S:
        return 1000
    if frequency == ConstantItem.FREQUENCY_1MIN:
        return 60000
    elif frequency == ConstantItem.FREQUENCY_3MIN:
        return 180000
    elif frequency == ConstantItem.FREQUENCY_5MIN:
        return 300000
    elif frequency == ConstantItem.FREQUENCY_15MIN:
        return 900000
    elif frequency == ConstantItem.FREQUENCY_30MIN:
        return 1800000
    elif frequency == ConstantItem.FREQUENCY_1HOUR:
        return 3600000
    elif frequency == ConstantItem.FREQUENCY_2HOUR:
        return 7200000
    elif frequency == ConstantItem.FREQUENCY_4HOUR:
        return 14400000
    elif frequency == ConstantItem.FREQUENCY_6HOUR:
        return 21600000
    elif frequency == ConstantItem.FREQUENCY_8HOUR:
        return 28800000
    elif frequency == ConstantItem.FREQUENCY_12HOUR:
        return 43200000
    elif frequency == ConstantItem.FREQUENCY_1D:
        return 86400000
    elif frequency == ConstantItem.FREQUENCY_1W:
        return 604800000
    else:
        raise BizException(BasicExceptionCode.ILLEGAL_PARAMETER)


if __name__ == '__main__':
    db: Session = next(deps.getDb())
    
    start_date = DateUtils.parse("2023-06-16")
    
    # 1.采集商品信息
    # spider = JqdataStockInstrumentSpider(db)
    # spider.crawl()

    # 2.采集K线
    # instruments = db.query(Instrument).filter(Instrument.marketType == MarketType.STOCK).all()
    # spier = DefaultJqdataStockBarSpider(db, ConstantItem.FREQUENCY_1D, DateUtils.parse("2005-01-01"), instruments)
    # spier.crawl()
    
    instruments = db.query(Instrument).filter(Instrument.id.in_([5188])).all()
    spier = DefaultJqdataStockBarSpider(db, ConstantJQDataFrequency.FREQUENCY_1D, start_date, instruments)
    spier.crawl()
    
    # instruments = db.query(Instrument).filter(Instrument.id.in_([1])).all()
    # spier = DefaultJqdataStockBarSpider(db, ConstantJQDataFrequency.FREQUENCY_60MIN, start_date, instruments)
    # spier.crawl()